from typing import TypedDict, Annotated, List
import operator
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import InMemorySaver  # 示例，生产环境可能用Redis/DB


# 1. 定义图的状态
class AgentState(TypedDict):
    messages: Annotated[List[BaseMessage], operator.add]
    user_input: str  # 用于存储用户在人机交互环节的输入


# 2. 定义AI节点 (Agent Node)
def call_llm(state: AgentState):
    messages = state['messages']
    # 假设这里有一个简单的LLM调用
    # 实际中你会使用 LangChain 的 LLM 实例
    ai_response_content = f"AI thinking based on: {messages[-1].content}. Do you want to continue or modify?"
    return {"messages": [AIMessage(content=ai_response_content)]}


# 3. 定义人机交互节点 (Human-in-the-Loop Node)
# 这个节点会检查是否需要人工干预
def human_intervene(state: AgentState):
    # 这里可以定义触发人工干预的逻辑，例如：
    # if "敏感词" in state['messages'][-1].content:
    #     return "request_human_review"
    # 或者总是要求人工干预特定步骤
    if "continue" not in state['user_input'].lower() and "modify" not in state['user_input'].lower():
        print(f"DEBUG: AI needs human input. Current state: {state['messages'][-1].content}")
        # LangGraph 的 interrupt 功能，用于暂停执行
        # 在Web API中，这个中断会被捕获，并通知用户
        return "human_input_required"
    else:
        print(f"DEBUG: Human input received: {state['user_input']}")
        return "continue_processing"


# 4. 构建LangGraph
workflow = StateGraph(AgentState)

workflow.add_node("llm", call_llm)
workflow.add_node("human_intervention_check", human_intervene)

workflow.set_entry_point("llm")

# 定义边
workflow.add_edge("llm", "human_intervention_check")

# 定义条件路由
workflow.add_conditional_edges(
    "human_intervention_check",
    lambda state: state["user_input"],  # 路由条件基于user_input
    {
        "human_input_required": END,  # 如果需要人工输入，则流程在这里“结束”（等待外部触发恢复）
        "continue_processing": "llm",  # 如果用户选择继续，则回到LLM（例如，处理用户的新输入或澄清）
    }
)

# 编译图，并使用一个检查点管理器
# InMemorySaver 适合本地测试，生产环境请使用 RedisCheckpointManager 或 SQLCheckpointManager
my_graph = workflow.compile(checkpointer=InMemorySaver())
